#%%
### This script was written by Eric Deal, 06.2022
### Script should be run with python 3
import numpy as np, pandas as pd, pathlib, seaborn as sns, matplotlib.pyplot as plt, scipy.special as ss
grains = ['dn', 'jo', 'mb', 'mk', 'pg', 'sf']
grains_new = {'dn': 'tg', 'jo': 'ng2', 'mb': 'ng3', 'mk': 'ng4', 'pg': 'ng5', 'sf': 'sf'}

fns = ['dragon glass-Table 1.csv',
        'Joel\'s gravel-Table 1.csv',
        'marble chips-Table 1.csv',
        'Michael\'s gravel-Table 1.csv',
        'pea gravel-Table 1.csv',
        'shell fragments-Table 1.csv']

grain_name = {'dragon glass-Table 1.csv': 'dn',
        'Joel\'s gravel-Table 1.csv': 'jo',
        'marble chips-Table 1.csv': 'mb',
        'Michael\'s gravel-Table 1.csv': 'mk',
        'pea gravel-Table 1.csv': 'pg',
        'shell fragments-Table 1.csv': 'sf'}

df_shape_taylor = {}
df_shape_taylor_std = {}
df_n = {}
for fn in fns:
    temp = pd.read_csv('observations/2022_taylor_shape_data/'+fn)
    df_n[grain_name[fn]] = temp.CSF.shape[0]
    df_shape_taylor[grain_name[fn]] = {'A': temp['a (long)'].mean(), 'B': temp['b (mid)'].mean(), 'C': temp['c (short)'].mean(), 'CSF': temp['CSF'].mean()}
    df_shape_taylor_std[grain_name[fn]] = {'A': temp['a (long)'].std(), 'B': temp['b (mid)'].std(), 'C': temp['c (short)'].std(), 'CSF': temp['CSF'].std()}

df_shape_taylor = pd.DataFrame(df_shape_taylor).T
df_shape_taylor_std = pd.DataFrame(df_shape_taylor_std).T
n = {grains_new[grain]: df_n[grain] for grain in grains}
A = {grains_new[grain]: df_shape_taylor.A[grain]/1000 for grain in grains}
dA = {grains_new[grain]:df_shape_taylor_std.A[grain]/1000/np.sqrt(df_n[grain]) for grain in grains}
B = {grains_new[grain]: df_shape_taylor.B[grain]/1000 for grain in grains}
dB = {grains_new[grain]: df_shape_taylor_std.B[grain]/1000/np.sqrt(df_n[grain]) for grain in grains}
C = {grains_new[grain]: df_shape_taylor.C[grain]/1000 for grain in grains}
dC = {grains_new[grain]: df_shape_taylor_std.C[grain]/1000/np.sqrt(df_n[grain]) for grain in grains}
CSF = {grains_new[grain]: C[grains_new[grain]]/np.sqrt(A[grains_new[grain]]*B[grains_new[grain]]) for grain in grains}
# dCSF = {grains_new[grain]: CSF[grains_new[grain]]*np.sqrt((dA[grains_new[grain]]/(2*A[grains_new[grain]]))**2 + (dB[grains_new[grain]]/(2*B[grains_new[grain]]))**2 + (dC[grains_new[grain]]/(C[grains_new[grain]]))**2) for grain in grains}
dCSF = {grains_new[grain]: df_shape_taylor_std.CSF[grain]/np.sqrt(df_n[grain]) for grain in grains}



df = pd.read_csv('1_var_density_2022.csv').set_index('Unnamed: 0').rename_axis(None)
dn = ((6/np.pi)*df.grain_vol)**(1/3)
ddn = dn*df.dgrain_vol/(3*df.grain_vol)

# combine all the data into a single pandas dataframe
df_save = pd.DataFrame({'A': A, 'dA': dA, 'B': B, 'dB': dB, 'C': C, 'dC': dC, 'dn': dn, 'ddn': ddn, 'CSF': CSF, 'dCSF': dCSF, 'n': n})
# save data to .csv file.
df_save.to_csv('2_var_shape_2022.csv')

# %%
